import asyncio
import os
import aiohttp
import requests
from urllib.parse import urlparse, unquote


class DownloadManager:
    def __init__(self, concurrency):
        self.download_tasks = []  # 存储下载任务的列表
        self.current_downloads = set()  # 当前正在下载的任务集合
        self.progress_callback = None  # 下载进度回调函数
        self.complete_callback = None  # 所有任务完成回调函数
        self.total_bytes_expected = 0  # 所有任务总体大小
        self.total_bytes_downloaded = 0  # 已下载的字节数
        self.semaphore = asyncio.Semaphore(concurrency)
        self.donwloaded = 0  # 已下载数量

    async def download(self, url, save_path, filename):
        if filename is None:
            filename = self.get_filename_from_url(url)

        if not filename:
            print(f"无法获取 {url} 的文件名")
            return

        save_path = os.path.join(save_path, filename)

        # 创建保存路径的目录
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        try:
            async with self.semaphore:
                async with aiohttp.ClientSession(
                    connector=aiohttp.TCPConnector(force_close=False)
                ) as session:
                    async with session.get(url) as response:
                        total_size = int(response.headers.get("Content-Length", 0))
                        self.total_bytes_expected += total_size  # 增加总体大小

                        bytes_downloaded = 0

                        with open(save_path, "wb") as f:
                            async for chunk in response.content.iter_any():
                                if chunk:
                                    f.write(chunk)
                                    bytes_downloaded += len(chunk)

                                    # 计算总体进度
                                    total_progress = (
                                        (
                                            self.total_bytes_downloaded
                                            / self.total_bytes_expected
                                        )
                                        * 100
                                        if self.total_bytes_expected > 0
                                        else 0
                                    )

                                    # # 单个文件进度
                                    # progress = (
                                    #     (bytes_downloaded / total_size) * 100
                                    #     if total_size > 0
                                    #     else 0
                                    # )

                                    if self.progress_callback1:
                                        self.progress_callback1(total_progress)

                                    self.total_bytes_downloaded += len(chunk)
        except Exception as e:
            print(f"下载 {url} 时发生其他错误: {str(e)}")
            raise  # 如果是其他错误，直接向上层抛出异常

        self.donwloaded += 1

        percentage = (self.donwloaded / len(self.download_tasks)) * 100
        percentage_rounded = round(percentage, 2)

        self.progress_callback("", str(percentage_rounded))
        # print(
        #     "当前进度" + str(len(self.download_tasks)) + "占" + str(percentage_rounded)
        # )

        # 完成下载后的回调
        # if self.complete_callback:
        #     self.complete_callback(filename)

    def get_filename_from_url(self, url):
        try:
            response = requests.head(url, allow_redirects=True)
            if "Content-Disposition" in response.headers:
                content_disposition = response.headers["Content-Disposition"]
                filename = content_disposition.split("filename=")[1]
                return unquote(filename.strip("\"'"))
            else:
                parsed_url = urlparse(url)
                filename = os.path.basename(parsed_url.path)
                return filename
        except Exception as e:
            print(f"获取文件名出错: {str(e)}")
            return None

    async def start_downloads(self):
        self.total_bytes_expected = 0  # 重置总体大小
        self.total_bytes_downloaded = 0  # 重置已下载大小

        for url, save_path, filename in self.download_tasks:
            download_coroutine = self.download(url, save_path, filename)
            self.current_downloads.add(download_coroutine)

        await asyncio.gather(*self.current_downloads)

        # 清空下载任务列表和当前下载集合
        self.download_tasks.clear()
        self.current_downloads.clear()

        print("所有下载任务已完成！")

        # 所有任务完成后的回调
        if self.complete_callback:
            self.complete_callback(None)  # 可以传递任何需要的参数

    def set_progress_callback(self, callback):
        self.progress_callback = callback

    def set_progress_callback1(self, callback):
        self.progress_callback1 = callback

    def set_complete_callback(self, callback):
        self.complete_callback = callback

    def add_download_task(self, url, save_path, filename=None):
        self.download_tasks.append((url, save_path, filename))


# 示例用法
# if __name__ == "__main__":
#     downloader = DownloadManager(6)  # 携带并发参数

#     # 自定义下载进度回调函数
#     def progress_callback(filename, progress):
#         print(f"{filename} 下载进度: {progress}%")

#     def progress_callback1(progress):
#         print(f"下载进度: {progress:.2f}%")

#     # 自定义所有任务完成回调函数
#     def complete_callback(result):
#         if result is None:
#             print("所有下载任务已完成！")

#     # 设置回调函数
#     downloader.set_progress_callback(progress_callback)
#     downloader.set_progress_callback1(progress_callback1)
#     downloader.set_complete_callback(complete_callback)

#     # 添加下载任务
#     downloader.add_download_task(
#         "https://wirelesscdn-download.xuexi.cn/publish/xuexi_android/latest/xuexi_android_10002068.apk",
#         "downloads",
#         "zzz.zip",
#     )
#     downloader.add_download_task(
#         "https://dldir1.qq.com/qqfile/qq/PCQQ9.7.17/QQ9.7.17.29225.exe",
#         "downloads",
#         "qq.apk",
#     )

#     # 异步开始下载
#     asyncio.run(downloader.start_downloads())
